require(Rniftilib)
require(bigmemory)
require(foreach)

source("base.R")
source("neuroimage.R")

# Default is to allow parellel computing if setup
# options(allowParellel=FALSE) if want otherwise
options(allowParallel=TRUE)

setClass(
    Class = "nifticor", 
    representation = representation(
        header = "list",
        mask_used = "vector",
        descriptorfile = "character"
    ),
    contains = "big.matrix"
)

# Questions?
# - do i need to remove an old instance of a file-backed big matrix or will it automatically overwrite

setGeneric('is.nifticor', 
    function(x) standardGeneric('is.nifticor')
)

setMethod('is.nifticor',
    signature(x='nifticor'),
    function(x) return(TRUE)
)

setMethod('is.nifticor',
    definition=function(x) return(FALSE)
)

nifticor <- function(x, output_prefix, m=NULL, blocksize=1000, overwrite=F, verbose=T) {
    # Check input args
    if (class(x) != "nifti_matrix")
        stop("input must be a nifti_matrix object")
    if (!is.character(output_prefix))
        stop("output_prefix must be character")
    
    # Mask
    if (!is.null(m))
        x <- mask(x, m)
    
    # Setup paths
    output_prefix <- rmext(abspath(output_prefix))
    backingpath <- dirname(output_prefix)
    backingfile <- sprintf("%s.bin", basename(output_prefix))
    descriptorfile <- sprintf("%s.desc", basename(output_prefix))
    nifticorfile <- sprintf("%s.nifticor", basename(output_prefix))
    
    # Check output exists
    # and possibly return old output
    if (file.exists(file.path(backingpath, backingfile))) {
        warning(sprintf("file-backed correlation matrix '%s/%s' already exists", backingpath, backingfile))
        if (!overwrite) {
            warning("Returning previous version")
            return(attach.nifticor(file.path(backingpath, nifticorfile)))
        }
    }
    
    # Create empty big matrix
    nVoxs <- nrow(x)
    nTimePts <- ncol(x)
    nBlocks <- ceiling(nVoxs/blocksize)
    cormat <- big.matrix(nVoxs, nVoxs, init=1, type="double", backingpath=backingpath, backingfile=backingfile, descriptorfile=descriptorfile)
    
    if (verbose) {
        pb <- create.progressbar(nBlocks)
        update.progressbar(pb, 0)
    }
    
    # Scale Time Series
    scalemat <- aperm(x) # Flip
    scalemat <- scale(scalemat)
    
    # Compute cross-product in chunks
    lapply(1:nBlocks, function(i) {
        blockstart <- (i-1)*blocksize + 1
        blockend <- min(nVoxs, blockstart + blocksize - 1)
        indices <- blockstart:blockend
        
        cormat[indices,] <- crossprod(scalemat[,indices], scalemat)/(nTimePts-1)
        
        if(verbose) update.progressbar(pb, i)
    })
    
    # New class
    cormat <- new("nifticor", address=cormat@address, header=x@header, mask_used=x@mask_used, descriptorfile=file.path(backingpath, descriptorfile))
    
    # Save info
    dput(list(header=cormat@header, mask_used=cormat@mask_used, descriptorfile=cormat@descriptorfile), file=file.path(backingpath, nifticorfile))
    
    if (verbose)
        end.progressbar(pb)
    
    return(cormat)
}
 
attach.nifticor <- function(fname) {
    if (!file.exists(fname))
        stop("input does not exist ", fname)
    nlist <- dget(fname)
    cormat <- attach.big.matrix(nlist$descriptorfile)
    cormat <- new("nifticor", address=cormat@address, header=nlist$header, mask_used=nlist$mask_used, descriptorfile=nlist$descriptorfile)
    return(cormat)
}

setMethod('coords',
    signature(x='nifticor'),
    function(x) .coords(x@header, x@mask_used)
)

batch.nifticor <- function(inputs, outputs, m=NULL, verbose=T, ...) {
    if (length(inputs) != length(outputs))
        stop("Different number of inputs than outputs")
    
    if (verbose)
        cat("Reading in data\n")
    
    inputs <- batch.nifti2mat(inputs, m, verbose=verbose)
    
    if (verbose)
        cat("Computing correlations\n")
    
    corfun <- function(i) nifticor(inputs[[i]], outputs[[i]], verbose=verbose, ...)
    
    if (getOption("allowParallel") && getDoParRegistered()) {
        corlist <- foreach(i=1:length(inputs), .packages=c("bigmemory", "Rniftilib")) %dopar% corfun(i)
    } else {
        corlist <- lapply(1:length(inputs), function(i) {
            if (verbose) cat(sprintf("\nRunning subject %03i", i), "\n");
            corfun
        })
    }
        
    gc(reset=T, verbose=F)
    
    return(corlist)
}

getnifticors <- function(xs, m=NULL) {
    m <- getmask(m)

    xs <- lapply(xs, function(x) {
        if (is.character(x))
            x <- attach.nifticor(paste(rmext(x), ".nifticor", sep=""))
        
        if (!is.nifticor(x))
            stop("unrecognized type ", class(x))
        
        if (!is.null(m))
            x <- mask(x, m)
        
        return(x)    
    })
    
    tmp1 <- sapply(xs, nrow)
    tmp2 <- sapply(xs, ncol)
    if (!all(tmp1==tmp1[1]) || !all(tmp2==tmp2[1]))
        stop("number of voxels must be same across all inputs")    
    
    return(xs)
}
